from typing import Callable, Dict, List, Optional, Union

from lagent.schema import AgentMessage


class Memory:
    """
    代理系统的基础内存管理类。
    用于存储和管理代理消息的历史记录。

    Attributes:
        memory (List[AgentMessage]): 存储代理消息的列表
        recent_n (Optional[int]): 获取最近消息时的默认数量限制
    """

    def __init__(self, recent_n=None) -> None:
        """
        初始化内存管理器。

        Args:
            recent_n (Optional[int]): 获取最近消息时的默认数量限制
        """
        self.memory: List[AgentMessage] = []
        self.recent_n = recent_n

    def get_memory(
        self,
        recent_n: Optional[int] = None,
        filter_func: Optional[Callable[[int, dict], bool]] = None,
    ) -> list:
        """
        获取存储的消息。

        Args:
            recent_n (Optional[int]): 要获取的最近消息数量，覆盖实例的默认值
            filter_func (Optional[Callable]): 用于过滤消息的函数，接收索引和消息作为参数

        Returns:
            list: 符合条件的消息列表
        """
        recent_n = recent_n or self.recent_n
        if recent_n is not None:
            memory = self.memory[-recent_n:]  # 获取最近的 n 条消息
        else:
            memory = self.memory  # 获取所有消息
        if filter_func is not None:
            memory = [m for i, m in enumerate(memory) if filter_func(i, m)]  # 应用过滤函数
        return memory

    def add(self, memories: Union[List[Dict], Dict, None]) -> None:
        """
        添加新的消息到内存中。

        Args:
            memories (Union[List[Dict], Dict, None]): 要添加的消息，可以是单个消息或消息列表
                                                    支持字符串、字典或 AgentMessage 对象
        """
        for memory in memories if isinstance(memories,
                                             (list, tuple)) else [memories]:
            if isinstance(memory, str):
                memory = AgentMessage(sender='user', content=memory)  # 将字符串转换为 AgentMessage
            if isinstance(memory, AgentMessage):
                self.memory.append(memory)

    def delete(self, index: Union[List, int]) -> None:
        """
        从内存中删除指定索引的消息。

        Args:
            index (Union[List, int]): 要删除的消息索引，可以是单个索引或索引列表
        """
        if isinstance(index, int):
            del self.memory[index]
        else:
            for i in index:
                del self.memory[i]

    def load(
        self,
        memories: Union[str, Dict, List],
        overwrite: bool = True,
    ) -> None:
        """
        从外部数据加载消息到内存中。

        Args:
            memories (Union[str, Dict, List]): 要加载的消息数据
            overwrite (bool): 是否覆盖现有的内存内容，默认为 True

        Raises:
            TypeError: 当提供的数据类型不支持时抛出
        """
        if overwrite:
            self.memory = []  # 如果 overwrite 为 True，清空现有内存
        if isinstance(memories, dict):
            self.memory.append(AgentMessage(**memories))
        elif isinstance(memories, list):
            for m in memories:
                self.memory.append(AgentMessage(**m))
        else:
            raise TypeError(f'{type(memories)} is not supported')

    def save(self) -> List[dict]:
        """
        将内存中的消息保存为可序列化的字典列表。

        Returns:
            List[dict]: 包含所有消息数据的字典列表
        """
        memory = []
        for m in self.memory:
            memory.append(m.model_dump())  # 将每个 AgentMessage 转换为字典
        return memory
